from langchain.text_splitter import CharacterTextSplitter
from tool import custom_print

"""
Split by character:
    官方文档：https://python.langchain.com/docs/modules/data_connection/document_transformers/text_splitters/character_text_splitter
    
    1.CharacterTextSplitter 可设置 separator、chunk_size 和 chunk_overlap 三个主要属性；
    2.CharacterTextSplitter 的 separator 属性只能设置一种字符进行分割

"""

with open('test2.txt', encoding='utf-8') as f:
    state_of_the_union = f.read()


# 自定义分割1: 固定 50 长度截取，不保留重叠字符
# fixed_size_text_splitter = CharacterTextSplitter(
#     separator="",
#     chunk_size=50,
#     chunk_overlap=0,
#     length_function=len,
# )
# fixed_size_texts = fixed_size_text_splitter.create_documents([state_of_the_union])
# custom_print.print_all(fixed_size_texts)


# 自定义分割2: 固定 50 长度截取，保留 10 个重叠字符
# fixed_size_chunk_overlap10_text_splitter = CharacterTextSplitter(
#     separator="",
#     chunk_size=50,
#     chunk_overlap=10,
#     length_function=len,
# )
# fixed_size_chunk_overlap10_text_texts = fixed_size_chunk_overlap10_text_splitter.create_documents([state_of_the_union])
# custom_print.print_all(fixed_size_chunk_overlap10_text_texts)


# 自定义分割3: 以单个换行符分割截取，最大长度为 50，保留 10 个重叠字符
split_by_wrap_chunk_overlap10_text_splitter = CharacterTextSplitter(
    separator="\n",
    chunk_size=50,
    chunk_overlap=10,
    length_function=len,
)
split_by_wrap_chunk_overlap10_texts = split_by_wrap_chunk_overlap10_text_splitter.create_documents([state_of_the_union])
custom_print.print_all(split_by_wrap_chunk_overlap10_texts)
